import pandas as pd
import numpy as np

import pandas as pd
import numpy as np

def compare_dataframes_safely(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    key_cols: list = ['date', 'code'],
    compare_cols: list = None,
    chunksize: int = 10000,
    tolerance: float = 1e-6
) -> pd.DataFrame:
    """
    内存安全的数据框对比方法 (支持多列主键和大数据)
    
    参数:
        df1, df2: 要对比的两个DataFrame
        key_cols: 唯一标识行的列 (必须存在于两个df中)
        compare_cols: 需要对比的数值列 (默认对比所有非主键列)
        chunksize: 分块处理的行数 (根据内存调整)
        tolerance: 浮点数对比容差
        
    返回:
        包含差异信息的DataFrame，格式:
        [key_cols..., 'diff_columns', 'value_df1', 'value_df2']
    """
    # ========== 预处理阶段 ==========
    # 1. 校验主键列存在
    for df in [df1, df2]:
        missing = set(key_cols) - set(df.columns)
        if missing:
            raise ValueError(f"主键列 {missing} 不存在于数据中")

    # 2. 日期列类型转换
    date_cols = [col for col in key_cols if 'date' in col.lower()]
    for col in date_cols:
        df1[col] = pd.to_datetime(df1[col])
        df2[col] = pd.to_datetime(df2[col])

    # 3. 确定对比列
    all_cols = list(df1.columns.intersection(df2.columns))
    compare_cols = compare_cols or [c for c in all_cols if c not in key_cols]
    
    # ========== 核心对比逻辑 ==========
    # 生成唯一键 (优化内存)
    key_series = df1[key_cols].apply(lambda x: '_'.join(x.astype(str)), axis=1)
    unique_keys = key_series.unique()
    
    # 准备结果容器
    diff_records = []
    
    # 分块处理
    for i in range(0, len(unique_keys), chunksize):
        chunk_keys = unique_keys[i:i+chunksize]
        
        # 提取当前块数据
        mask1 = key_series.isin(chunk_keys)
        df1_chunk = df1.loc[mask1, key_cols + compare_cols].reset_index(drop=True)
        df2_chunk = df2[df2[key_cols].apply(
            lambda x: '_'.join(x.astype(str)), axis=1).isin(chunk_keys)
        ].reset_index(drop=True)
        
        # 合并对比
        merged = pd.merge(
            df1_chunk, 
            df2_chunk, 
            on=key_cols, 
            suffixes=('_df1', '_df2'),
            how='outer', 
            indicator=True
        )
        
        # 处理差异
        for _, row in merged.iterrows():
            record = {col: row[col] for col in key_cols}
            
            # 1. 处理缺失行
            if row['_merge'] != 'both':
                diff_cols = ['缺失行']
                record.update({
                    'diff_columns': ', '.join(diff_cols),
                    'value_df1': '存在' if row['_merge'] == 'left_only' else '缺失',
                    'value_df2': '存在' if row['_merge'] == 'right_only' else '缺失'
                })
                diff_records.append(record)
                continue
                
            # 2. 数值对比
            diff_cols = []
            for col in compare_cols:
                v1 = row[f'{col}_df1']
                v2 = row[f'{col}_df2']
                
                # 处理空值
                if pd.isna(v1) or pd.isna(v2):
                    if pd.isna(v1) != pd.isna(v2):
                        diff_cols.append(col)
                    continue
                
                # 数值对比
                if pd.api.types.is_numeric_dtype(type(v1)):
                    if not np.isclose(v1, v2, atol=tolerance, equal_nan=True):
                        diff_cols.append(col)
                else:
                    if v1 != v2:
                        diff_cols.append(col)
            
            if diff_cols:
                for col in compare_cols:
                    if col in diff_cols:
                        record[f'{col}_df1'] = row[f'{col}_df1']
                        record[f'{col}_df2'] = row[f'{col}_df2']
                record['diff_columns'] = ', '.join(diff_cols)
                diff_records.append(record)
    
    # ========== 结果整理 ==========
    if not diff_records:
        return pd.DataFrame(columns=key_cols + ['diff_columns'])
    
    result = pd.DataFrame(diff_records)
    result = result.sort_values(key_cols).reset_index(drop=True)
    return result

# # ========== 使用示例 ==========
# if __name__ == "__main__":
#     # 测试数据
#     df1 = pd.DataFrame({
#         'date': ['2023-01-01', '2023-01-01', '2023-01-02'],
#         'code': ['A', 'B', 'A'],
#         'close': [100.0, 200.0, 105.0],
#         'volume': [10000, 20000, 15000]
#     })
    
#     df2 = pd.DataFrame({
#         'date': ['2023-01-01', '2023-01-02'],
#         'code': ['A', 'A'],
#         'close': [100.1, 105.0],  # 微小差异
#         'volume': [10000, 15000]
#     })
    
#     # 执行对比
#     diff_df = compare_dataframes_safely(
#         df1, 
#         df2,
#         key_cols=['date', 'code'],
#         chunksize=5000,
#         tolerance=0.2
#     )
    
#     print("差异结果:")
#     print(diff_df)

# diff_result = compare_dataframes(single_df, single_df2)
# diff_result